from zai import ZhipuAiClient

class AdaptiveRAGPipeline:
    def __init__(self, api_key, model_name):
        self.model = model_name
        self.client = ZhipuAiClient(api_key=api_key)  # 请填写您自己的 API Key

    def query(self, question: str) -> str:
        """
        简化的 Adaptive RAG 实现，调用大模型返回答案。
        可以在此扩展：检索 -> 生成 -> 自适应调整。
        """
        response = self.client.chat.completions.create(
            model=self.model,
            messages=[
                {"role": "user", "content": question}
            ],
            thinking={
                "type": "enabled",    # 启用深度思考模式
            },
            max_tokens=65536,          # 最大输出 tokens
            temperature=1.0           # 控制输出的随机性
        )
        # 获取完整回复
        msg = response.choices[0].message # type: ignore
        print(msg)
        return msg.content # type: ignore
